package com.maggie.demo.scgsentinel.handler;

import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.BlockRequestHandler;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.Ordered;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.server.ServerResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.HashSet;
import java.util.Set;

import static org.springframework.web.reactive.function.BodyInserters.fromValue;

@Component
@Slf4j
public class SentinelBlockRequestHandler implements BlockRequestHandler, Ordered {
  
  public static final Set<String> FLOW_LIMIT_API = new HashSet<>(32);
  
  static {
    FLOW_LIMIT_API.add("/test/api/tab/test");
  }

  public static final String CODE_SYSTEM_BUSY = "11004";

  public static final String MSG_SYSTEM_BUSY = "网络开小差了,请稍后重试.";

  @Override
  public int getOrder() {
    return LOWEST_PRECEDENCE;
  }
  
  @Autowired
  private WebClient.Builder webClientBuilder;
  
  @Override
  public Mono<ServerResponse> handleRequest(ServerWebExchange exchange, Throwable ex) {

    String requestUrl = exchange.getRequest().getPath().toString();
    Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
    if (FLOW_LIMIT_API.contains(requestUrl)) {
      return handleRequestDetail(exchange,ex);
    }

    // JSON result by default.
    return ServerResponse.status(HttpStatus.OK)
        .contentType(MediaType.APPLICATION_JSON)
        .body(fromValue(buildErrorResult(ex)));
  }

  private ErrorResult buildErrorResult(Throwable ex) {
    return new ErrorResult(CODE_SYSTEM_BUSY, MSG_SYSTEM_BUSY);
  }

  private static class ErrorResult {

    private final String status;
    private final String message;

    public String getStatus() {
      return status;
    }

    ErrorResult(String status, String message) {
      this.status = status;
      this.message = message;
    }

    public String getMessage() {
      return message;
    }
  }
  
  
  public Mono<ServerResponse> handleRequestDetail(ServerWebExchange exchange, Throwable ex) {
    String requestUrl = exchange.getRequest().getURI().getRawPath();
    System.out.println("requestUrl = " + requestUrl);
    Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
    // post请求的参数需要自己写filter 获取,这里只是获取的get请求参数
    MultiValueMap<String, String> getParams = exchange.getRequest().getQueryParams();
    
    // 转发到第三方接口
    StringBuilder newUri = new StringBuilder("http://")
            .append(route.getId())
            .append(requestUrl.substring("/test".length())+"/sentinel");
    
    log.info("限流uri: {}, query参数: {}, 降级至: uri {}", requestUrl, getParams, newUri.toString());
    //
    return webClientBuilder.build()
            .post() //
            .uri(newUri.toString())
            // 传递 request header
            .headers(newHeaders -> newHeaders.putAll(exchange.getRequest().getHeaders()))
            .bodyValue(getParams)
            .exchangeToMono(response -> {
              return response.bodyToMono(String.class)
                      .defaultIfEmpty("")
                      .flatMap(body -> {
                        return ServerResponse.status(response.statusCode())
                                // 传递 response header,避免cookie在网关层丢失
                                .headers(it -> {
                                  it.addAll(response.headers().asHttpHeaders());
                                })
                                .bodyValue(body);
                      });
            });
  }
}
